import colorsys
import os
import warnings

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from utils.augmentations import letterbox
import datetime
from utils.box_utils import letterbox_image, ssd_correct_boxes
from nets_student.ssd_student_layers import Detect
warnings.filterwarnings("ignore")

MEANS = (104, 117, 123)

class SSD(object):
    _defaults = {
        "classes_path": 'model_data/new_classes.txt',
        "input_shape": (512, 512, 3),
        "cuda": False,
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, opt, **kwargs):
        self.opt = opt
        self.__dict__.update(self._defaults)
        self.class_names = self._get_class()
        self.generate()

    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def generate(self):
        self.num_classes = len(self.class_names) + 1
        device = torch.device('cuda' if self.opt.cuda else 'cpu')
        model = torch.load(self.opt.target_weights, map_location=device)
        model.phase = 'detection'
        self.net = model.eval()
        if self.opt.cuda:
            self.net = torch.nn.DataParallel(self.net)
            cudnn.benchmark = True
            self.net = self.net.cuda()

        self.detect = Detect(self.num_classes, 0, 200, self.opt.conf_thres, self.opt.iou_thres)
        print('model, anchors, and classes loaded.')
        # 画框设置不同的颜色
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))

    def detect_image(self, image):
        with torch.no_grad():
            image0 = image  # copy image
            image_shape = image.shape[:2]  # get original image shape
            image = letterbox(image, self.opt.input_shape)  # 这里的letterbox采用的yolov5的

            image = image[0] - MEANS
            image = image.transpose((2, 0, 1))[::-1]  # HWC->CHW
            image = np.ascontiguousarray(image)

            images = torch.from_numpy(image)  # numpy to tensor
            if self.cuda:
                images = images.cuda()
            else:
                images = images.float()
            if len(image.shape) == 3:
                photo = images[None]  # expand for batch dim
            preds = self.net(photo)
            loc = preds[0]
            conf = preds[1]
            priors = preds[2]
            preds = self.detect.forward(loc[0], nn.Softmax(dim=-1).forward(conf[0]), priors)
            top_conf = []
            top_label = []
            top_bboxes = []
            for i in range(preds.size(1)):
                j = 0
                while preds[0, i, j, 0] >= self.opt.conf_thres:  # preds[...,0] 第0个维度是conf
                    score = preds[0, i, j, 0]
                    label_name = self.class_names[i - 1]
                    pt = (preds[0, i, j, 1:]).detach().numpy()
                    coords = [pt[0], pt[1], pt[2], pt[3]]  # 这里的坐标是相对于网络输入大小图像的而不是原图
                    top_conf.append(score)
                    top_label.append(label_name)
                    top_bboxes.append(coords)
                    j = j + 1

        # 如果不存在满足门限的预测框，直接返回原图
        if len(top_conf) <= 0:
            return image0

        top_conf = np.array(top_conf)
        top_label = np.array(top_label)
        top_bboxes = np.array(top_bboxes)
        top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:, 0], -1), \
                                                 np.expand_dims(top_bboxes[:, 1], -1), \
                                                 np.expand_dims(top_bboxes[:, 2], -1), \
                                                 np.expand_dims(top_bboxes[:, 3], -1)
        boxes = ssd_correct_boxes(top_ymin, top_xmin, top_ymax, top_xmax,
                                  np.array([self.opt.input_shape, self.opt.input_shape]), image_shape)

        thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.opt.input_shape, 1)

        for i, c in enumerate(top_label):
            predicted_class = c
            score = top_conf[i]
            # 获取当前时间
            current_time = datetime.datetime.now()
            formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
            # 将记录写入data/data.txt文件
            with open('output/data/data.txt', 'a') as file:
                file.write(f"{formatted_time} - 已检测到目标\n")
            top, left, bottom, right = boxes[i]  # y1,x1,y2,x2
            top = top - 5
            left = left - 5
            bottom = bottom + 5
            right = right + 5

            top = max(0, np.floor(top + 0.5).astype('int32'))  # y1
            left = max(0, np.floor(left + 0.5).astype('int32'))  # x1
            bottom = min(image0.shape[0], np.floor(bottom).astype('int32'))  # 原图的H和boxes右下角y比较防止越界
            right = min(image0.shape[1], np.floor(right).astype('int32'))
            axis = [top, left, bottom, right]

            label = '{} {:.2f}'.format(predicted_class, score)
            cv2.putText(image0, label, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0))
            #print(label, top, left, bottom, right)
            for i in range(thickness):
                cv2.rectangle(image0,  (left, top), (right, bottom), (255, 0, 0), 1)
        return image0, axis


